#include <path_searching/astar.h>
#include <sstream>

using namespace std;
using namespace Eigen;

namespace fast_planner {
    Astar::~Astar() {
        for (size_t i = 0; i < allocate_num_; i++) {
            delete path_node_pool_[i];
        }
    }

    int Astar::search(const Eigen::Vector3d &start_pt, const Eigen::Vector3d &end_pt, bool dynamic, double time_start) {
        /* ---------- initialize ---------- */
        NodePtr cur_node = path_node_pool_[0];
        cur_node->parent = nullptr;
        cur_node->position = start_pt;
        cur_node->index = posToIndex(start_pt);
        cur_node->g_score = 0.0;
        cur_node->f_score = lambda_heu_ * getEuclHeu(cur_node->position, end_pt);
        cur_node->node_state = IN_OPEN_SET;

        Eigen::Vector3i end_index = posToIndex(end_pt);

        open_set_.push(cur_node);
        use_node_num_ += 1;
        expanded_nodes_.insert(cur_node->index, cur_node);
        NodePtr terminate_node = nullptr;

        /* ---------- search loop ---------- */
        while (!open_set_.empty()) {
            /* ---------- get lowest f_score node ---------- */
            cur_node = open_set_.top();

            const int end_range = 2;
            bool reach_end = abs(cur_node->index(0) - end_index(0)) <= end_range &&
                             abs(cur_node->index(1) - end_index(1)) <= end_range &&
                             abs(cur_node->index(2) - end_index(2)) <= end_range;
            if (reach_end) {
                terminate_node = cur_node;
                retrievePath(terminate_node);
                has_path_ = true;
                return REACH_END;
            }

            /* ---------- pop node and add to close set ---------- */
            open_set_.pop();
            cur_node->node_state = IN_CLOSE_SET;
            iter_num_ += 1;
            Eigen::Vector3d cur_pos = cur_node->position;
            Eigen::Vector3d pro_pos;
            Eigen::Vector3d d_pos;

            /* ---------- expansion loop ---------- */
            for (double dx = -resolution_; dx <= resolution_ + 1e-3; dx += resolution_)
                for (double dy = -resolution_; dy <= resolution_ + 1e-3; dy += resolution_)
                    for (double dz = -resolution_; dz <= resolution_ + 1e-3; dz += resolution_) {
                        d_pos << dx, dy, dz;
                        if (d_pos.norm() < 1e-3) continue;
                        pro_pos = cur_pos + d_pos;
                        /* ---------- check if in feasible space ---------- */
                        if (pro_pos(0) <= origin_(0) || pro_pos(0) >= map_size_3d_(0) || pro_pos(1) <= origin_(1) ||
                            pro_pos(1) >= map_size_3d_(1) || pro_pos(2) <= origin_(2) ||
                            pro_pos(2) >= map_size_3d_(2)) {
                            continue;
                        }
                        /* not in close set */
                        Eigen::Vector3i pro_id = posToIndex(pro_pos);
                        NodePtr pro_node = expanded_nodes_.find(pro_id);
                        if (pro_node != nullptr && pro_node->node_state == IN_CLOSE_SET) {
                            continue;
                        }
                        /* collision free */
                        double dist = edt_environment_->sdf_map_->getDistance(pro_pos);
                        if (dist <= margin_) {
                            continue;
                        }
                        /* ---------- compute cost ---------- */
                        double tmp_g_score, tmp_f_score;
                        tmp_g_score = d_pos.squaredNorm() + cur_node->g_score;
                        tmp_f_score = tmp_g_score + lambda_heu_ * getEuclHeu(pro_pos, end_pt);

                        if (pro_node == nullptr) {
                            pro_node = path_node_pool_[use_node_num_];
                            pro_node->index = pro_id;
                            pro_node->position = pro_pos;
                            pro_node->f_score = tmp_f_score;
                            pro_node->g_score = tmp_g_score;
                            pro_node->parent = cur_node;
                            pro_node->node_state = IN_OPEN_SET;

                            open_set_.push(pro_node);
                            expanded_nodes_.insert(pro_id, pro_node);

                            use_node_num_ += 1;
                            if (use_node_num_ == allocate_num_) {
                                cout << "run out of memory." << endl;
                                return NO_PATH;
                            }
                        } else if (pro_node->node_state == IN_OPEN_SET) {
                            if (tmp_g_score < pro_node->g_score) {
                                // pro_node->discretized_state = pro_id;
                                pro_node->position = pro_pos;
                                pro_node->f_score = tmp_f_score;
                                pro_node->g_score = tmp_g_score;
                                pro_node->parent = cur_node;
                            }
                        } else {
                            cout << "error type in searching: " << pro_node->node_state << endl;
                        }
                    }
        }

        /* ---------- open set empty, no path ---------- */
        cout << "open set empty, no path!" << endl;
        cout << "use node num: " << use_node_num_ << endl;
        cout << "iter num: " << iter_num_ << endl;
        return NO_PATH;
    }

    void Astar::setParam(ros::NodeHandle &nh) {
        nh.param("astar/resolution_astar", resolution_, -1.0);
        nh.param("astar/time_resolution", time_resolution_, -1.0);
        nh.param("astar/lambda_heu", lambda_heu_, -1.0);
        nh.param("astar/margin", margin_, -1.0);
        nh.param("astar/allocate_num", allocate_num_, -1);
        tie_breaker_ = 1.0 + 1.0 / 10000;
        cout << "margin:" << margin_ << endl;
    }

    void Astar::retrievePath(NodePtr end_node) {
        NodePtr cur_node = end_node;
        path_nodes_.push_back(cur_node);
        while (cur_node->parent != nullptr) {
            cur_node = cur_node->parent;
            path_nodes_.push_back(cur_node);
        }
        reverse(path_nodes_.begin(), path_nodes_.end());
    }

    std::vector<Eigen::Vector3d> Astar::getPath() {
        vector<Eigen::Vector3d> path;
        for (auto &path_node: path_nodes_) {
            path.push_back(path_node->position);
        }
        return path;
    }

    double Astar::getDiagHeu(Eigen::Vector3d x1, Eigen::Vector3d x2) const {
        double dx = fabs(x1(0) - x2(0));
        double dy = fabs(x1(1) - x2(1));
        double dz = fabs(x1(2) - x2(2));
        double h;
        int diag = min(min(dx, dy), dz);
        dx -= diag;
        dy -= diag;
        dz -= diag;

        if (dx < 1e-4) {
            h = 1.0 * sqrt(3.0) * diag + sqrt(2.0) * min(dy, dz) + 1.0 * abs(dy - dz);
        }
        if (dy < 1e-4) {
            h = 1.0 * sqrt(3.0) * diag + sqrt(2.0) * min(dx, dz) + 1.0 * abs(dx - dz);
        }
        if (dz < 1e-4) {
            h = 1.0 * sqrt(3.0) * diag + sqrt(2.0) * min(dx, dy) + 1.0 * abs(dx - dy);
        }
        return tie_breaker_ * h;
    }

    double Astar::getManhHeu(Eigen::Vector3d x1, Eigen::Vector3d x2) {
        double dx = fabs(x1(0) - x2(0));
        double dy = fabs(x1(1) - x2(1));
        double dz = fabs(x1(2) - x2(2));
        return tie_breaker_ * (dx + dy + dz);
    }

    double Astar::getEuclHeu(const Eigen::Vector3d &x1, const Eigen::Vector3d &x2) const {
        return tie_breaker_ * (x2 - x1).norm();
    }

    void Astar::init() {
        /* ---------- map params ---------- */
        this->inv_resolution_ = 1.0 / resolution_;
        inv_time_resolution_ = 1.0 / time_resolution_;
        edt_environment_->sdf_map_->getRegion(origin_, map_size_3d_);
        cout << "astar origin_: " << origin_.transpose() << endl;
        cout << "astar map size: " << map_size_3d_.transpose() << endl;

        path_node_pool_.resize(allocate_num_);
        for (size_t i = 0; i < allocate_num_; i++) {
            path_node_pool_[i] = new Node;
        }
        use_node_num_ = 0;
        iter_num_ = 0;
    }

    void Astar::setEnvironment(const EDTEnvironment::Ptr &env) {
        this->edt_environment_ = env;
    }

    void Astar::reset() {
        expanded_nodes_.clear();
        path_nodes_.clear();

        std::priority_queue<NodePtr, std::vector<NodePtr>, NodeComparator0> empty_queue;
        open_set_.swap(empty_queue);
        for (size_t i = 0; i < use_node_num_; i++) {
            NodePtr node = path_node_pool_[i];
            node->parent = nullptr;
            node->node_state = NOT_EXPAND;
        }
        use_node_num_ = 0;
        iter_num_ = 0;
    }

    std::vector<Eigen::Vector3d> Astar::getVisited() {
        vector<Eigen::Vector3d> visited;
        for (size_t i = 0; i < use_node_num_; ++i)
            visited.push_back(path_node_pool_[i]->position);
        return visited;
    }

    Eigen::Vector3i Astar::posToIndex(const Eigen::Vector3d &pt) {
        Vector3i idx = ((pt - origin_) * inv_resolution_).array().floor().cast<int>();
        return idx;
    }

    int Astar::timeToIndex(double time) const {
        int idx = floor((time - time_origin_) * inv_time_resolution_);
    }

}  // namespace fast_planner
